package kzgX

import (
	"fmt"
	"kzg-ipa-study/tools"

	kbls "github.com/kilic/bls12-381"
	"github.com/protolambda/go-kzg/bls"
)

// v is big-endian bytes
func FrFromByte32(dst *bls.Fr, v []byte) {
	(*kbls.Fr)(dst).RedFromBytes(v[:])
}

// NegFr 取反
func NegFr(in bls.Fr) bls.Fr {
	var tmp bls.Fr
	bls.SubModFr(&tmp, &bls.ZERO, &in)
	return tmp
}

// Long polynomial division for two polynomials in coefficient form
func polyLongDiv(dividend []bls.Fr, divisor []bls.Fr) []bls.Fr {
	a := make([]bls.Fr, len(dividend), len(dividend))
	for i := 0; i < len(a); i++ {
		bls.CopyFr(&a[i], &dividend[i])
	}
	aPos := len(a) - 1
	bPos := len(divisor) - 1
	diff := aPos - bPos
	out := make([]bls.Fr, diff+1, diff+1)
	for diff >= 0 {
		quot := &out[diff]

		bls.DivModFr(quot, &a[aPos], &divisor[bPos])
		var tmp, tmp2 bls.Fr
		for i := bPos; i >= 0; i-- {
			// In steps: a[diff + i] -= b[i] * quot
			// tmp =  b[i] * quot
			bls.MulModFr(&tmp, quot, &divisor[i])
			// tmp2 = a[diff + i] - tmp
			bls.SubModFr(&tmp2, &a[diff+i], &tmp)
			// a[diff + i] = tmp2
			bls.CopyFr(&a[diff+i], &tmp2)
		}
		aPos -= 1
		diff -= 1
	}
	return out
}

func polynomialMul(a, b []bls.Fr) []bls.Fr {
	r := arrayOfZeroes(len(a) + len(b) - 1)

	for i := 0; i < len(a); i++ {
		for j := 0; j < len(b); j++ {
			// r[i+j] = fAdd(r[i+j], fMul(a[i], b[j]))
			var tmp bls.Fr
			bls.MulModFr(&tmp, &a[i], &b[j])
			bls.AddModFr(&r[i+j], &r[i+j], &tmp)
		}
	}
	return r
}

func arrayOfZeroes(n int) []bls.Fr {
	r := make([]bls.Fr, n)
	for i := 0; i < n; i++ {
		r[i] = bls.ZERO
	}
	return r[:]
}

// LagrangeInterpolation implements the Lagrange interpolation:
// https://en.wikipedia.org/wiki/Lagrange_polynomial
func LagrangeInterpolation(x, y []bls.Fr) ([]bls.Fr, error) {
	if len(x) != len(y) {
		return nil, fmt.Errorf("len(x)!=len(y): %d, %d", len(x), len(y))
	}

	p := arrayOfZeroes(len(x))
	k := len(x)

	for j := 0; j < k; j++ {
		// jPol is the Lagrange basis polynomial for each point
		var jPol []bls.Fr
		for m := 0; m < k; m++ {
			// if x[m] == x[j] {
			if m == j {
				continue
			}
			// numerator & denominator of the current iteration
			num := []bls.Fr{NegFr(x[m]), bls.ONE} // (x^1 - x_m)
			// den := fSub(x[j], x[m])               // x_j-x_m
			var den bls.Fr
			bls.SubModFr(&den, &x[j], &x[m])
			mPol := polynomialDivByConstant(num, den)
			if len(jPol) == 0 {
				// first j iteration
				jPol = mPol
				continue
			}
			jPol = polynomialMul(jPol, mPol)
		}
		p = polynomialAdd(p, polynomialMulByConstant(jPol, y[j]))
	}

	return p, nil
}

func polynomialAdd(a, b []bls.Fr) []bls.Fr {
	r := arrayOfZeroes(tools.Max(len(a), len(b)))
	for i := 0; i < len(a); i++ {
		// r[i] = fAdd(r[i], a[i])
		bls.AddModFr(&r[i], &r[i], &a[i])
	}
	for i := 0; i < len(b); i++ {
		// r[i] = fAdd(r[i], b[i])
		bls.AddModFr(&r[i], &r[i], &b[i])
	}
	return r
}

func polynomialMulByConstant(a []bls.Fr, c bls.Fr) []bls.Fr {
	for i := 0; i < len(a); i++ {
		// a[i] = fMul(a[i], c)
		bls.MulModFr(&a[i], &a[i], &c)
	}
	return a
}

func polynomialDivByConstant(a []bls.Fr, c bls.Fr) []bls.Fr {
	for i := 0; i < len(a); i++ {
		// a[i] = fDiv(a[i], c)
		bls.DivModFr(&a[i], &a[i], &c)
	}
	return a
}

func polynomialSub(a, b []bls.Fr) []bls.Fr {
	r := arrayOfZeroes(tools.Max(len(a), len(b)))
	for i := 0; i < len(a); i++ {
		// r[i] = fAdd(r[i], a[i])
		bls.AddModFr(&r[i], &r[i], &a[i])
	}
	for i := 0; i < len(b); i++ {
		// r[i] = fSub(r[i], b[i])
		bls.SubModFr(&r[i], &r[i], &b[i])
	}
	return r
}

func polynomialDiv(a, b []bls.Fr) ([]bls.Fr, []bls.Fr) {
	// https://en.wikipedia.org/wiki/Division_algorithm
	r := arrayOfZeroes(len(a) - len(b) + 1)
	rem := a
	for len(rem) >= len(b) {
		// l := fDiv(rem[len(rem)-1], b[len(b)-1])
		var l bls.Fr
		bls.DivModFr(&l, &rem[len(rem)-1], &b[len(b)-1])
		pos := len(rem) - len(b)
		r[pos] = l
		aux := arrayOfZeroes(pos)
		aux1 := append(aux, l)
		aux2 := polynomialSub(rem, polynomialMul(b, aux1))
		rem = aux2[:len(aux2)-1]
	}
	return r, rem
}

func compareFrArray(a, b []bls.Fr) bool {
	if len(a) != len(b) {
		return false
	}
	for i := 0; i < len(a); i++ {
		// if a[i] != b[i] {
		if !bls.EqualFr(&a[i], &b[i]) {
			return false
		}
	}
	return true
}

func evaluateG1(g1s []bls.G1Point, p []bls.Fr) *bls.G1Point {
	return bls.LinCombG1(g1s, p)
}

func evaluateG2(g2s []bls.G2Point, p []bls.Fr) *bls.G2Point {
	return LinCombG2(g2s, p)
}

func LinCombG2(numbers []bls.G2Point, factors []bls.Fr) *bls.G2Point {
	if len(numbers) != len(factors) {
		panic("got LinCombG2 numbers/factors length mismatch")
	}
	var out bls.G2Point
	tmpG2s := make([]*kbls.PointG2, len(numbers), len(numbers))
	for i := 0; i < len(numbers); i++ {
		tmpG2s[i] = (*kbls.PointG2)(&numbers[i])
	}
	tmpFrs := make([]*kbls.Fr, len(factors), len(factors))
	for i := 0; i < len(factors); i++ {
		// copy, since we have to change from mont-red form to regular again, and don't want to mutate the input.
		v := *(*kbls.Fr)(&factors[i])
		v.FromRed()
		tmpFrs[i] = &v
	}
	_, err := kbls.NewG2().MultiExp((*kbls.PointG2)(&out), tmpG2s, tmpFrs)
	if err != nil {
		panic(err)
	}
	return &out
}
